import torch
import torch.nn as nn

# from mmdet.core import bbox2result, bbox2roi, build_assigner, build_sampler
from mmdet.models.builder import DETECTORS, build_backbone, build_head, build_neck
from .obb_base import OBBBaseDetector
from .obb_test_mixins import RotateAugRPNTestMixin
import torch.nn.functional as F
import numpy as np
from mmdet.core import add_prefix
import warnings


def resize(input,
           size=None,
           scale_factor=None,
           mode='nearest',
           align_corners=None,
           warning=True):
    if warning:
        if size is not None and align_corners:
            input_h, input_w = tuple(int(x) for x in input.shape[2:])
            output_h, output_w = tuple(int(x) for x in size)
            if output_h > input_h or output_w > output_h:
                if ((output_h > 1 and output_w > 1 and input_h > 1
                     and input_w > 1) and (output_h - 1) % (input_h - 1)
                        and (output_w - 1) % (input_w - 1)):
                    warnings.warn(
                        f'When align_corners={align_corners}, '
                        'the output would more aligned if '
                        f'input size {(input_h, input_w)} is `x+1` and '
                        f'out size {(output_h, output_w)} is `nx+1`')
    return F.interpolate(input, size, scale_factor, mode, align_corners)






@DETECTORS.register_module()
class OBBTwoStageDetector_interact(OBBBaseDetector, RotateAugRPNTestMixin):
    """Base class for two-stage detectors.

    Two-stage detectors typically consisting of a region proposal network and a
    task-specific regression head.
    """

    def __init__(self,
                 backbone,
                 neck=None,
                 interact_module = None,
                 cls_head=None,
                 rpn_head=None,
                 roi_head=None,
                 seg_head=None,
                 auxiliary_seg_head=None,
                 train_cfg=None,
                 test_cfg=None,
                 pretrained=None):
        super(OBBTwoStageDetector_interact, self).__init__()
        self.backbone = build_backbone(backbone)

        if neck is not None:
            self.neck = build_neck(neck)

        if interact_module is not None:
            self.interact_module = build_neck(interact_module)

        
        if rpn_head is not None:
            rpn_train_cfg = train_cfg.rpn if train_cfg is not None else None
            rpn_head_ = rpn_head.copy()
            rpn_head_.update(train_cfg=rpn_train_cfg, test_cfg=test_cfg.rpn)
            self.rpn_head = build_head(rpn_head_)

        if roi_head is not None:
            # update train and test cfg here for now
            # TODO: refactor assigner & sampler
            rcnn_train_cfg = train_cfg.rcnn if train_cfg is not None else None
            roi_head.update(train_cfg=rcnn_train_cfg)
            roi_head.update(test_cfg=test_cfg.rcnn)
            self.roi_head = build_head(roi_head)

        if seg_head is not None:
            # update train and test cfg here for now
            # TODO: refactor assigner & sampler
            seg_train_cfg = train_cfg.seg if train_cfg is not None else None
            seg_head.update(train_cfg=seg_train_cfg)
            seg_head.update(test_cfg=test_cfg.seg)
            self.seg_head = build_head(seg_head)
            self.align_corners = self.seg_head.align_corners


        if auxiliary_seg_head is not None:
            # auxiliary_seg_train_cfg = train_cfg.seg if train_cfg is not None else None
            # auxiliary_seg_head.update(train_cfg=auxiliary_seg_train_cfg)
            # auxiliary_seg_head.update(test_cfg=test_cfg.seg)
            self.auxiliary_seg_head = build_head(auxiliary_seg_head)
        
        if cls_head is not None:
            # cls_train_cfg = None
            # cls_head_ = cls_head.copy()
            # cls_head_.update(train_cfg=cls_train_cfg, test_cfg=cls_train_cfg)
            self.cls_head = build_head(cls_head)    

        self.train_cfg = train_cfg
        self.test_cfg = test_cfg

        self.init_weights(pretrained=pretrained)

    @property
    def with_rpn(self):
        """bool: whether the detector has RPN"""
        return hasattr(self, 'rpn_head') and self.rpn_head is not None

    @property
    def with_roi_head(self):
        """bool: whether the detector has a RoI head"""
        return hasattr(self, 'roi_head') and self.roi_head is not None
    
    @property
    def with_interact_module(self):
        """bool: whether the detector has seg_head"""
        return hasattr(self, 'interact_module') and self.interact_module is not None
    

    @property
    def with_seg_head(self):
        """bool: whether the detector has seg_head"""
        return hasattr(self, 'seg_head') and self.seg_head is not None

    @property
    def with_aux_seg_head(self):
        """bool: whether the detector has a RoI head"""
        return hasattr(self, 'auxiliary_seg_head') and self.auxiliary_seg_head is not None
    
    @property
    def with_cls_head(self):
        """bool: whether the detector has a cls_head"""
        return hasattr(self, 'cls_head') and self.cls_head is not None
    


    def init_weights(self, pretrained=None):
        """Initialize the weights in detector

        Args:
            pretrained (str, optional): Path to pre-trained weights.
                Defaults to None.
        """
        super(OBBTwoStageDetector_interact, self).init_weights(pretrained)
        self.backbone.init_weights(pretrained=pretrained)
        if self.with_neck:
            if isinstance(self.neck, nn.Sequential):
                for m in self.neck:
                    m.init_weights()
            else:
                self.neck.init_weights()
        if self.with_rpn:
            self.rpn_head.init_weights()
        if self.with_roi_head:
            self.roi_head.init_weights(pretrained)
        if self.with_interact_module:
            self.interact_module.init_weights()

    def extract_feat(self, img):
        """Directly extract features from the backbone+neck
        """
        x = self.backbone(img)
        if self.with_neck:
            x = self.neck(x)
        if self.with_interact_module:
            # print(x)
            ss_feat, rd_feat, cls_feat = self.interact_module(x)
            # print(len(x),ss_feat.shape,rd_feat.shape,cls_feat.shape)
            # 5 torch.Size([1, 64, 256, 256]) torch.Size([1, 64, 256, 256]) torch.Size([1, 64, 256, 256])
            # for i in range(len(ss_feat)):
            #     print(f'ss_feat[{i}]: {ss_feat[i].shape}')
            
            # for i in range(len(rd_feat)):
            #     print(f'rd_feat[{i}]: {rd_feat[i].shape}')

            # print(cls_feat.shape)




            # x[0]: torch.Size([1, 256, 256, 256])
            # x[1]: torch.Size([1, 256, 128, 128])
            # x[2]: torch.Size([1, 256, 64, 64])
            # x[3]: torch.Size([1, 256, 32, 32])
            # x[4]: torch.Size([1, 256, 16, 16])
                
            return ss_feat, rd_feat, cls_feat
        else:
            return x

    def forward_dummy(self, img):
        """Used for computing network flops.

        See `mmdetection/tools/get_flops.py`
        """
        outs = ()
        # backbone
        x = self.extract_feat(img)
        # rpn
        proposal_type = 'hbb'
        if self.with_rpn:
            proposal_type = getattr(self.rpn_head, 'bbox_type', 'hbb')
            rpn_outs = self.rpn_head(x)
            outs = outs + (rpn_outs, )

        if proposal_type == 'hbb':
            proposals = torch.randn(1000, 4).to(img.device)
        elif proposal_type == 'obb':
            proposals = torch.randn(1000, 5).to(img.device)
        else:
            # poly proposals need to be generated in roi_head
            proposals = None
        # roi_head
        roi_outs = self.roi_head.forward_dummy(x, proposals)
        outs = outs + (roi_outs, )
        return outs

    def forward_train(self,
                      img,
                      img_metas,
                      gt_bboxes,
                      gt_obboxes,
                      gt_labels,
                      gt_bboxes_ignore=None,
                      gt_obboxes_ignore=None,
                      proposals=None,
                      **kwargs):
        """
        Args:
            img (Tensor): of shape (N, C, H, W) encoding input images.
                Typically these should be mean centered and std scaled.

            img_metas (list[dict]): list of image info dict where each dict
                has: 'img_shape', 'scale_factor', 'flip', and may also contain
                'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
                For details on the values of these keys see
                `mmdet/datasets/pipelines/formatting.py:Collect`.

            gt_bboxes (list[Tensor]): Ground truth bboxes for each image with
                shape (num_gts, 4) in [tl_x, tl_y, br_x, br_y] format.

            gt_labels (list[Tensor]): class indices corresponding to each box

            gt_bboxes_ignore (None | list[Tensor]): specify which bounding
                boxes can be ignored when computing the loss.

            gt_masks (None | Tensor) : true segmentation masks for each box
                used if the architecture supports a segmentation task.

            proposals : override rpn proposals with custom proposals. Use when
                `with_rpn` is False.

        Returns:
            dict[str, Tensor]: a dictionary of loss components
        """
        ss_feat, rd_feat, cls_feat  = self.extract_feat(img)



        losses = dict()

        # RPN forward and loss
        if self.with_rpn:
            proposal_type = getattr(self.rpn_head, 'bbox_type', 'hbb')
            target_bboxes = gt_bboxes if proposal_type == 'hbb' else gt_obboxes
            target_bboxes_ignore = gt_bboxes_ignore if proposal_type == 'hbb' \
                    else gt_obboxes_ignore

            proposal_cfg = self.train_cfg.get('rpn_proposal',
                                              self.test_cfg.rpn)
            rpn_losses, proposal_list = self.rpn_head.forward_train(
                rd_feat,
                img_metas,
                target_bboxes,
                gt_labels=None,
                gt_bboxes_ignore=target_bboxes_ignore,
                proposal_cfg=proposal_cfg)
            losses.update(rpn_losses)
        else:
            proposal_list = proposals

        if self.with_roi_head:
            roi_losses = self.roi_head.forward_train(rd_feat, img_metas, proposal_list,
                                                    gt_bboxes, gt_obboxes, gt_labels,
                                                    gt_bboxes_ignore, gt_obboxes_ignore,
                                                    **kwargs)
            losses.update(roi_losses)
        
        if self.with_seg_head:
            seg_losses = self.seg_head.forward_train(ss_feat, img_metas, train_cfg=None, **kwargs)
            losses.update(seg_losses)
        
        if self.with_aux_seg_head:
            seg_losses_aux = self.auxiliary_seg_head.forward_train(ss_feat, img_metas, train_cfg=None, **kwargs)
            losses.update(add_prefix(seg_losses_aux, 'aux'))

        if self.with_cls_head:
            cls_losses = self.cls_head.forward_train(cls_feat, img_metas, **kwargs)
            losses.update(cls_losses)

        return losses
    

    async def async_simple_test(self,
                                img,
                                img_meta,
                                proposals=None,
                                rescale=False):
        """Async test without augmentation."""
        assert self.with_bbox, 'Bbox head must be implemented.'
        x = self.extract_feat(img)

        if proposals is None:
            proposal_list = await self.rpn_head.async_simple_test_rpn(
                x, img_meta)
        else:
            proposal_list = proposals

        return await self.roi_head.async_simple_test(
            x, proposal_list, img_meta, rescale=rescale)

    def simple_test(self, img, img_metas, proposals=None, rescale=False):
        """Test without augmentation."""
        # assert self.with_bbox, 'Bbox head must be implemented.'
        outs = {}
        ss_feat, rd_feat, cls_feat  = self.extract_feat(img)

        
        if self.with_rpn:
            if proposals is None:
                proposal_list = self.rpn_head.simple_test_rpn(rd_feat, img_metas)
            else:
                proposal_list = proposals
        
        if self.with_roi_head:
            roi_out = self.roi_head.simple_test(rd_feat, proposal_list, img_metas, rescale=rescale)
            outs['det_out'] = roi_out

        if self.with_seg_head:
            outs['seg_out'] = self.seg_inference(img, ss_feat, img_metas, rescale)

        if self.with_cls_head:
            outs['cls_out'] = np.array(self.cls_head.forward_test(cls_feat).cpu())

        return outs

    def aug_test(self, imgs, img_metas, rescale=False):
        """Test with augmentations.

        If rescale is False, then returned bboxes and masks will fit the scale
        of imgs[0].
        """
        # recompute feats to save memory
        x = self.extract_feats(imgs)
        proposal_list = self.rotate_aug_test_rpn(x, img_metas)
        return self.roi_head.aug_test(
            x, proposal_list, img_metas, rescale=rescale)



    def slide_inference(self, img, img_meta, rescale):
        """Inference by sliding-window with overlap.

        If h_crop > h_img or w_crop > w_img, the small patch will be used to
        decode without padding.
        """

        h_stride, w_stride = self.test_cfg['seg'].stride
        h_crop, w_crop = self.test_cfg['seg'].crop_size
        batch_size, _, h_img, w_img = img.size()
        num_classes = self.seg_head.num_classes
        h_grids = max(h_img - h_crop + h_stride - 1, 0) // h_stride + 1
        w_grids = max(w_img - w_crop + w_stride - 1, 0) // w_stride + 1
        preds = img.new_zeros((batch_size, num_classes, h_img, w_img))
        count_mat = img.new_zeros((batch_size, 1, h_img, w_img))
        for h_idx in range(h_grids):
            for w_idx in range(w_grids):
                y1 = h_idx * h_stride
                x1 = w_idx * w_stride
                y2 = min(y1 + h_crop, h_img)
                x2 = min(x1 + w_crop, w_img)
                y1 = max(y2 - h_crop, 0)
                x1 = max(x2 - w_crop, 0)
                crop_img = img[:, :, y1:y2, x1:x2]

                # crop_seg_logit = self.encode_decode(crop_img, img_meta)
                x = self.backbone(crop_img)
                if self.with_neck:
                    x = self.neck(x)
                crop_seg_logit = self.seg_head.forward_test(x, img_meta, test_cfg=None)
                crop_seg_logit = resize(
                    input=crop_seg_logit,
                    size=crop_img.shape[2:],
                    mode='bilinear',
                    align_corners=self.align_corners)
                preds += F.pad(crop_seg_logit,
                               (int(x1), int(preds.shape[3] - x2), int(y1),
                                int(preds.shape[2] - y2)))

                count_mat[:, :, y1:y2, x1:x2] += 1
        assert (count_mat == 0).sum() == 0
        if torch.onnx.is_in_onnx_export():
            # cast count_mat to constant while exporting to ONNX
            count_mat = torch.from_numpy(
                count_mat.cpu().detach().numpy()).to(device=img.device)
        preds = preds / count_mat
        if rescale:
            preds = resize(
                preds,
                size=img_meta[0]['ori_shape'][:2],
                mode='bilinear',
                align_corners=False,
                warning=False)
        return preds

    def whole_inference(self, x, img_meta, rescale):
        """Inference with full image."""

        seg_logit = self.seg_head.forward_test(x, img_meta, test_cfg=None)
        if rescale:
            # support dynamic shape for onnx
            if torch.onnx.is_in_onnx_export():
                size = img.shape[2:]
            else:
                size = img_meta[0]['ori_shape'][:2]
            seg_logit = resize(
                seg_logit,
                size=size,
                mode='bilinear',
                align_corners=False,
                warning=False)

        return seg_logit


    def seg_inference(self, img, x, img_meta, rescale):
        """Inference with slide/whole style.

        Args:
            img (Tensor): The input image of shape (N, 3, H, W).
            img_meta (dict): Image info dict where each dict has: 'img_shape',
                'scale_factor', 'flip', and may also contain
                'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
                For details on the values of these keys see
                `mmseg/datasets/pipelines/formatting.py:Collect`.
            rescale (bool): Whether rescale back to original shape.

        Returns:
            Tensor: The output segmentation map.
        """

        assert self.test_cfg['seg'].mode in ['slide', 'whole']
        ori_shape = img_meta[0]['ori_shape']
        assert all(_['ori_shape'] == ori_shape for _ in img_meta)
        if self.test_cfg['seg'].mode == 'slide':
            seg_logit = self.slide_inference(img, img_meta, rescale)
        else:
            seg_logit = self.whole_inference(x, img_meta, rescale)
        output = F.softmax(seg_logit, dim=1)
        # flip = img_meta[0]['flip']
        # if flip:
        #     flip_direction = img_meta[0]['flip_direction']
        #     assert flip_direction in ['horizontal', 'vertical']
        #     if flip_direction == 'horizontal':
        #         output = output.flip(dims=(3, ))
        #     elif flip_direction == 'vertical':
        #         output = output.flip(dims=(2, ))

        # return output
        return np.array(torch.argmax(output, dim=1).squeeze().cpu())